import matplotlib.pyplot as plt
import numpy as np

# 假设你收集了不同GPU数量下的速度数据（samples/sec）
gpus = [2, 4, 8]
allreduce_throughput = [2.19, 2.19*1.2, 2.19*1.2*1.2]  # AllReduce下的速度
netsenseml_throughput =  [2.12, 2.12*1.2, 2.12*1.2*1.2]  # NetSenseML下的速度
static_compression_throughput = [2.12, 2.12*1.2, 2.12*1.2*1.2] # 静态压缩下的速度

# 创建图表
fig, ax = plt.subplots(figsize=(8, 6))

# 绘制三种方法下的训练速度对比图
width = 0.2  # 柱状图宽度
x = np.arange(len(gpus))  # GPU数量的横坐标

# 绘制AllReduce的柱状图
rects1 = ax.bar(x - width, allreduce_throughput, width, label='AllReduce', color='skyblue', hatch='//')

# 绘制NetSenseML的柱状图
rects2 = ax.bar(x, netsenseml_throughput, width, label='NetSenseML', color='orange', hatch='\\\\')

# 绘制静态压缩的柱状图
rects3 = ax.bar(x + width, static_compression_throughput, width, label='Static Compression', color='green', hatch='--')

# 设置标签和标题并增加字体大小
ax.set_xlabel('# of GPUs', fontsize=14)
ax.set_ylabel('Training Throughput (Samples/sec)', fontsize=14)
ax.set_title('GPT-2 Training Throughput Comparison', fontsize=16)
ax.set_xticks(x)
ax.set_xticklabels(gpus, fontsize=12)

# 设置图例字体大小
ax.legend(fontsize=12)

# 设置坐标轴刻度字体大小
ax.tick_params(axis='both', which='major', labelsize=12)

# 展示图表
plt.tight_layout()
plt.savefig('gpt2_training_throughput.pdf', format='pdf', bbox_inches='tight')

plt.show()